import os

from sqlalchemy.orm import sessionmaker
from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession

from ..utils.alembic_wrapper import AlembicMigrationDriver


class AsyncAlembicDbEngine:
    def __init__(self, db_url: str, alembic_db_url: str, **url_args):
        self.db_url = db_url
        self.alembic_db_url = alembic_db_url
        self._url_args = url_args
        self._path = self.db_url.format(**url_args)
        self._engine = create_async_engine(self._path)
        self._session = sessionmaker(self._engine, class_=AsyncSession, expire_on_commit=False)  # type: ignore

    @property
    def session(self):
        return self._session

    def create(self):
        driver = AlembicMigrationDriver(
            "mnagent.db.models.metadata",
            self.alembic_db_url.format(**self._url_args),
            alembic_prj_path=os.getcwd()
        )
        driver.update()

    def close(self):
        self._session.close_all()

    async def _query(self, sql):
        async with self.session() as session:
            return await session.execute(sql)

    async def _add(self, obj):
        async with self.session() as session:
            async with session.begin():
                session.add(obj)

    async def _update(self, sql):
        async with self.session() as session:
            async with session.begin():
                await session.execute(sql)

    update = _update
    add = _add
    query = _query
